package lk;

import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;


// 无限集中的最小数字
// 一个无限集, 初始包含所有的正整数
// popSmallest 方法，从 无限集 中移除并返回一个最小整数
// addBack(num) 方法，如果 无限集 中没有这个数，则添加进 无限集 中

class SmallestInfiniteSet {
    PriorityQueue<Integer> queue;   // 优先级队列，方便移除并返回 无限集 中的最小数
    int max;    // 记录 无限集 中 移除掉的最大数

    // 构造方法，进行初始化，初始状态 无限集 中，最小的数为 1, max = 0
    public SmallestInfiniteSet() {
        queue = new PriorityQueue<>();
        queue.offer(1);
        max = 0;
    }

    // 移除并返回最小的数
    public int popSmallest() {
        int min = queue.poll(); // 从优先级队列中取走最小的数

        if (queue.isEmpty()) {
            // 如果队列为空，说明是第一次移除这个数（也是目前移除的最大的数）
            // 更新max，并向队列中添加 无限集 中下一个最小的数
            max = min;
            queue.offer(min + 1);
        }
        return min;
    }

    public void addBack(int num) {
        // 如果这个数还在无限集中，就不必添加了
        if (queue.contains(num)) {
            return;
        }
        if (num > max) {
            return;
        }
        queue.offer(num);
    }
}

// 最大子序列的分数
// nums1 中选取 k 个数, 这 k 个数之和, 为 sum 1
// 这 k 个数的下标, 对应 nums2 中的 k 个数, nums2 的 k 个数的最小值为 y
// 求最大值 sum1 * y
public class D231008 {
    public long maxScore(int[] nums1, int[] nums2, int k) {
        int n = nums2.length;

        Integer[] indexs = new Integer[n];
        for (int i = 0; i < n; i++) {
            indexs[i] = i;
        }

        // index 中存储了下标, 根据 nums2 中元素的值, 对下标进行降序排序
        Arrays.sort(indexs, (o1, o2) -> nums2[o2] - nums2[o1]);
        // lambda 表达式(省略了参数类型、return、花括号等)，变量捕获，该变量的值不能修改

        // 结果
        long res = 0L;
        // 小根堆
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();
        // 从 nums1 中选取的子序列的和
        long sum1 = 0;

        // 遍历 indexs, 此时按照遍历的顺序 indexs 中的下标对应的 nums2 的元素值是递减的
        for (int i : indexs) {
            int x = nums1[i];
            int y = nums2[i];

            minHeap.offer(x);
            sum1 += x;

            if (minHeap.size() > k) {
                // size > k
                // 超过了 k 个元素, 移除之前选取的最小的元素
                sum1 -= minHeap.poll();
            }

            if (minHeap.size() == k) {
                // 刚好 k 个元素, 更新最大值
                // y 一直是当前的最小值
                long cur = sum1 * y;
                res = Math.max(res, cur);
            }
        }

        return res;
    }
}
